import os
import tkinter as tk
from tkinter import ttk, filedialog
import queue

import cv2
import matplotlib.pyplot as plt
import torch
from ultralytics import YOLO

# 加载 YOLOv11 Pose 模型
device = "cuda" if torch.cuda.is_available() else "cpu"
model = YOLO("./weights/yolo11x-shuttlecock.pt").to(device)

# 设置全局字体（任选一种）
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 黑体（Windows 系统自带）

# 解决负号显示异常
plt.rcParams["axes.unicode_minus"] = False

class ShuttlecockAnalysisApp:
    def __init__(self, root):
        self.root = root
        self.root.title("毽球动作骨骼关键点检测系统")
        self.root.geometry("600x400")

        # 视频文件路径
        self.video_path = tk.StringVar()
        # 动作名称
        self.motion_name = tk.StringVar()

        # 创建 UI
        self.create_widgets()

        # 图表初始化标志
        self.chart_initialized = False

        # 数据队列
        self.data_queue = queue.Queue()

        # 定时器任务
        self.update_chart_task()

    def create_widgets(self):
        # 主框架
        main_frame = ttk.Frame(self.root, padding="20")
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 视频选择部分
        video_frame = ttk.LabelFrame(main_frame, text="视频选择", padding="10")
        video_frame.pack(fill=tk.X, pady=5)

        ttk.Label(video_frame, text="视频路径:").grid(row=0, column=0, sticky=tk.W)
        ttk.Entry(video_frame, textvariable=self.video_path, width=50).grid(row=0, column=1, padx=5)
        ttk.Button(video_frame, text="浏览...", command=self.browse_video).grid(row=0, column=2)

        # 动作选择部分
        motion_frame = ttk.LabelFrame(main_frame, text="动作选择", padding="10")
        motion_frame.pack(fill=tk.X, pady=5)

        motion_options = [
            "左脚单踢-脚内侧踢球",
            "左脚单踢-脚外侧踢球",
            "左脚单踢-脚背踢球",
            "左脚单踢-腿部触球",
            "右脚单踢-脚内侧踢球",
            "右脚单踢-脚外侧踢球",
            "右脚单踢-脚背踢球",
            "右脚单踢-腿部触球",
            "双脚交替-脚内侧踢球",
            "双脚交替-脚背踢球",
            "双脚交替-腿部触球"
        ]

        ttk.Label(motion_frame, text="选择动作:").grid(row=0, column=0, sticky=tk.W)
        motion_combobox = ttk.Combobox(motion_frame, textvariable=self.motion_name,
                                       values=motion_options, state="readonly")
        motion_combobox.grid(row=0, column=1, padx=5, sticky=tk.W)
        motion_combobox.current(0)

        # 控制按钮
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=10)

        ttk.Button(button_frame, text="开始分析", command=self.start_analysis).pack(side=tk.LEFT, padx=5)
        ttk.Button(button_frame, text="退出", command=self.root.quit).pack(side=tk.RIGHT, padx=5)

    def browse_video(self):
        # 打开文件选择对话框，选择视频文件
        filepath = filedialog.askopenfilename(
            title="选择视频文件",
            filetypes=[("视频文件", "*.mp4 *.avi *.mov"), ("所有文件", "*.*")]
        )
        if filepath:
            self.video_path.set(filepath)

    def start_analysis(self):
        # 开始分析
        video_src = self.video_path.get()
        motion_name = self.motion_name.get()

        print(f"开始分析: {motion_name}")

        if not video_src:
            tk.messagebox.showerror("错误", "请先选择视频文件")
            return

        if not motion_name:
            tk.messagebox.showerror("错误", "请选择动作名称")
            return

        if not os.path.exists(video_src):
            tk.messagebox.showerror("错误", f"视频文件不存在: {video_src}")
            return

        # 在新线程中运行分析，避免阻塞 GUI
        import threading
        analysis_thread = threading.Thread(
            target=self._run_analysis,
            args=(video_src, motion_name),
            daemon=True
        )
        analysis_thread.start()

        # 在主线程中初始化图表
        if not self.chart_initialized:
            self.init_chart(self.get_keypoint_config(motion_name))
            plt.show(block=False)

    def _run_analysis(self, video_src, motion_name):
        # 在子线程中运行的分析函数
        vc = cv2.VideoCapture(video_src)
        if not vc.isOpened():
            print(f"错误：无法打开视频 {video_src}")
            return
        
        config = self.get_keypoint_config(motion_name)

        frame_count = 0

        while True:
            success, frame = vc.read()
            if not success:
                break

            frame_count += 1

            # 处理帧
            processed_frame, x1, y1, x2, y2 = self.process_frame(frame,  config)

            if processed_frame is None:
                break

            # 调整显示大小
            height, width = processed_frame.shape[:2]
            display_frame = cv2.resize(processed_frame, (width // 2, height // 2))

            # 显示状态
            status_text = f"Analying: {frame_count} frame"
            cv2.putText(display_frame, status_text, (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

            cv2.imshow('Processed Frame', display_frame)

            # 将数据放入队列
            self.data_queue.put((x1, y1, x2, y2, frame_count))

            # 按 'q' 键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        vc.release()
        cv2.destroyAllWindows()
        print(f"分析完成：{motion_name}")

    def update_chart_task(self):
        # 主线程定时任务，用于更新图表
        try:
            while True:
                # 从队列中获取数据
                x1, y1, x2, y2, frame_count = self.data_queue.get_nowait()
                self.update_chart(x1, y1, x2, y2, frame_count)
        except queue.Empty:
            pass

        # 每 100ms 检查一次队列
        self.root.after(100, self.update_chart_task)

    def init_chart(self, config):
        # 初始化图表
        self.fig, self.ax = plt.subplots(2, 2, figsize=(8, 4))
        plt.subplots_adjust(hspace=0.5)
        self.lines = []
        self.data_containers = {
            'relative_x1': [],
            'relative_y1': [],
            'relative_x2': [],
            'relative_y2': []
        }
        self.time_points = []

        titles = []
        for c in config.get('points'):
            titles.append(f'骨骼点 {c[0]} 与骨骼点 {c[1]} 在 X 方向上的差值')
            titles.append(f'骨骼点 {c[0]} 与骨骼点 {c[1]} 在 Y 方向上的差值')

        for i in range(4):
            row, col = i // 2, i % 2
            line, = self.ax[row, col].plot([], [], 'b-')
            if i < len(titles):
                self.ax[row, col].set_title(titles[i])
            self.ax[row, col].set_xlabel('帧')
            self.ax[row, col].grid(True)
            self.lines.append(line)

        plt.tight_layout()
        plt.ion()  # 开启交互模式
        plt.pause(0.1)  # 稍作暂停以初始化窗口
        self.chart_initialized = True

    def update_chart(self, relative_x1, relative_y1, relative_x2, relative_y2, frame_count):
        # 更新图表数据
        if not self.chart_initialized:
            return

        # 添加新数据
        self.data_containers['relative_x1'].append(relative_x1)
        self.data_containers['relative_y1'].append(relative_y1)
        self.data_containers['relative_x2'].append(relative_x2)
        self.data_containers['relative_y2'].append(relative_y2)
        self.time_points.append(frame_count)

        # 更新每个子图
        for i, key in enumerate(['relative_x1', 'relative_y1', 'relative_x2', 'relative_y2']):
            row, col = i // 2, i % 2
            self.lines[i].set_data(self.time_points, self.data_containers[key])
            self.ax[row, col].relim()
            self.ax[row, col].autoscale_view()

        # 强制重新绘制图表
        self.fig.canvas.draw()
        plt.pause(0.001)  # 稍作暂停以更新图表

    def process_frame(self, frame,  config):
        # 处理每一帧
        results = model(frame, verbose=False)[0]

        # 初始化相对坐标
        relative_x1, relative_y1 = 0, 0
        relative_x2, relative_y2 = 0, 0

        if results.keypoints is None:
            return frame, relative_x1, relative_y1, relative_x2, relative_y2

        kpts = results.keypoints.xy[0].cpu().numpy().astype(int)

        # 需要绘制的关键点索引
        desired_indices = [5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
        desired_pairs = [
            (5, 6), (6, 8), (8, 10), (5, 7),
            (11, 12), (12, 14), (14, 16), (13, 15),
            (11, 13), (7, 9), (6, 12), (5, 11)
        ]

        # 绘制关键点
        for idx in desired_indices:
            if idx < len(kpts):
                x, y = kpts[idx]
                cv2.circle(frame, (x, y), 6, (0, 255, 0), -1)

        # 绘制连接线
        for pair in desired_pairs:
            kp1, kp2 = pair
            if kp1 < len(kpts) and kp2 < len(kpts):
                x1, y1 = kpts[kp1]
                x2, y2 = kpts[kp2]
                cv2.line(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)

        # 处理关键点组（根据配置过滤）
        for group_idx, (x_idx, y_idx) in enumerate(config["points"], 1):
            if len(kpts) > max(x_idx, y_idx):
                x_val, y_val = kpts[x_idx]
                x_ref, y_ref = kpts[y_idx]

                # 计算相对坐标
                relative_x = x_val - x_ref
                relative_y = y_val - y_ref

                # 根据组索引存储数据
                if group_idx == 1:
                    relative_x1, relative_y1 = relative_x, relative_y
                elif group_idx == 2:
                    relative_x2, relative_y2 = relative_x, relative_y

        return frame, relative_x1, relative_y1, relative_x2, relative_y2

    def get_keypoint_config(self, motion_name: str):
        # 根据动作名称返回关键点配置
        configs = {
            # 左脚单踢动作
            "左脚单踢-脚内侧踢球": {"points": [(15, 14)]},
            "左脚单踢-脚外侧踢球": {"points": [(15, 11)]},
            "左脚单踢-脚背踢球": {"points": [(15, 13)]},
            "左脚单踢-腿部触球": {"points": [(13, 11)]},

            # 右脚单踢动作
            "右脚单踢-脚内侧踢球": {"points": [(16, 13)]},
            "右脚单踢-脚外侧踢球": {"points": [(16, 12)]},
            "右脚单踢-脚背踢球": {"points": [(16, 14)]},
            "右脚单踢-腿部触球": {"points": [(14, 12)]},

            # 双脚交替动作
            "双脚交替-脚内侧踢球": {"points": [(13, 14), (11, 12)]},
            "双脚交替-脚背踢球": {"points": [(16, 14), (15, 13)]},
            "双脚交替-腿部触球": {"points": [(14, 12), (13, 11)]},
        }
        return configs.get(motion_name, {"points": [(16, 13)]})  # 默认为右脚单踢


if __name__ == "__main__":
    root = tk.Tk()
    app = ShuttlecockAnalysisApp(root)
    root.mainloop()